from typing import NamedTuple

import numpy as np
from numpy.typing import ArrayLike
from zkl_aiutils_datasets import DelegatingDataset, MappedDataset

from .sliding_window import SlidingWindowDataset


class ArraySlidingWindowSample(NamedTuple):
    data: np.ndarray
    head: int
    tail: int


class ArraySlidingWindowDataset(DelegatingDataset[ArraySlidingWindowSample]):
    NoPadding = object()

    def __init__(self, array: np.ndarray, win_size: int, win_stride: int, pad_item: ArrayLike = NoPadding):
        if pad_item is self.NoPadding:
            org_len = None
        else:
            pad_size = win_size - len(array)
            if pad_size < 0:
                pad_size = pad_size % win_stride
            if pad_size > 0:
                org_len = len(array)
                array_pad = np.asarray(pad_item, dtype=array.dtype)
                array_pad = np.broadcast_to(array_pad, (pad_size,) + array.shape[1:])
                array = np.concatenate([array, array_pad], axis=0)
            else:
                org_len = None

        dataset = SlidingWindowDataset(array, win_size, win_stride)
        if org_len is None:
            dataset = MappedDataset(dataset, lambda sample: ArraySlidingWindowSample(*sample))
        else:
            dataset = MappedDataset(dataset, lambda sample: ArraySlidingWindowSample(
                data=sample.data, head=sample.head, tail=min(sample.tail, org_len)))

        super().__init__(dataset)
